import mxnet as mx
import numpy as np
import json


tmp_file = '/home/alpha/tmp/tmp.rec'
record = mx.recordio.MXRecordIO(tmp_file, 'w')
for i in range(5):
    record.write(('record_%d' % i).encode('utf-8'))
record.close()

record = mx.recordio.MXRecordIO(tmp_file, 'r')
while True:
    item = record.read()
    print(item)
    if not item:
        break
record.close()


idx_tmp_file = '/home/alpha/tmp/tmp.idx'
rec_tmp_file = '/home/alpha/tmp/tmp.rec'
record = mx.recordio.MXIndexedRecordIO(idx_tmp_file, rec_tmp_file, 'w')
for i in range(5):
    record.write_idx(i, ('record_%d' % i).encode('utf-8'))
record.close()

record = mx.recordio.MXIndexedRecordIO(idx_tmp_file, rec_tmp_file, 'r')
print(record.read_idx(3))
print(record.keys)
record.close()


# pack
data = json.dumps({
    'aa': [1,2,3]
}).encode('utf-8')

label1 = 1.0
header1 = mx.recordio.IRHeader(flag=0, label=label1, id=1, id2=0)
s1 = mx.recordio.pack(header1, data)
print('float label:', repr(s1))
label2 = [1.0, 2.0, 3.0]
header2 = mx.recordio.IRHeader(flag=0, label=label2, id=2, id2=0)
s2 = mx.recordio.pack(header2, data)
print('array label:', repr(s2))


print(*mx.recordio.unpack(s1))
print(*mx.recordio.unpack(s2))

